import os
import shutil

'''提取关键帧并命名'''
# 定义路径
rgb_folder = "E:\data\data_0905/room/rgb/"
depth_folder = "E:\data\data_0905/room/depth/"
keyframe_trajectory_file = "E:\data\data_0905/KeyFrameTrajectory2_name.txt"  # keyframetrajectory.txt文件路径
associate_file = "E:\data\data_0905/room/associate.txt"  # associate.txt文件路径

output_rgb_folder = "E:\change_data\c_data_0905key2/rgb"# 提取出的RGB图片存放路径
output_depth_folder = "E:\change_data\c_data_0905key2/depth"  # 提取出的Depth图片存放路径
# 创建输出文件夹
os.makedirs(output_rgb_folder, exist_ok=True)
os.makedirs(output_depth_folder, exist_ok=True)

# 读取keyframetrajectory.txt文件，获取关键帧的RGB图片名
keyframe_rgb_images = set()
with open(keyframe_trajectory_file, 'r') as file:
    for line in file:
        if line.strip():  # 跳过空行
            parts = line.strip().split()
            rgb_image_name = parts[0] + ".png"
            keyframe_rgb_images.add(rgb_image_name)

# 创建字典以存储RGB到Depth的映射
rgb_to_depth_mapping = {}

# 读取associate.txt文件，根据keyframe_rgb_images中记录的RGB图片名提取对应的Depth图片
with open(associate_file, 'r') as file:
    for line in file:
        if line.strip():  # 跳过空行
            parts = line.strip().split()
            rgb_image_name = parts[1].split('/')[-1]
            depth_image_path = parts[3]
            if rgb_image_name in keyframe_rgb_images:
                rgb_to_depth_mapping[rgb_image_name] = depth_image_path

# 按顺序重命名并复制RGB和Depth图片
index = 0
for rgb_image_name in sorted(keyframe_rgb_images):
    new_image_name = f"{index:06d}.png"

    # 复制并重命名RGB图片
    source_rgb_path = os.path.join(rgb_folder, rgb_image_name)
    target_rgb_path = os.path.join(output_rgb_folder, new_image_name)
    if os.path.exists(source_rgb_path):
        shutil.copy(source_rgb_path, target_rgb_path)
    else:
        print(f"Warning: {source_rgb_path} does not exist.")

    # 复制并重命名Depth图片
    if rgb_image_name in rgb_to_depth_mapping:
        depth_image_name = os.path.basename(rgb_to_depth_mapping[rgb_image_name])
        source_depth_path = os.path.join(depth_folder, depth_image_name)
        target_depth_path = os.path.join(output_depth_folder, new_image_name)
        if os.path.exists(source_depth_path):
            shutil.copy(source_depth_path, target_depth_path)
        else:
            print(f"Warning: {source_depth_path} does not exist.")

    index += 1

print("Processing complete.")
